'''
Created on Feb 1, 2010

@author: mkiyer
'''

import collections
import logging

import numpy as np
import matplotlib.pyplot as plt
import pysam

def depth_of_coverage(pileup_iterator, hist_dict):
    prev_pos = 0
    for pileupcolumn in pileup_iterator:
        pos = pileupcolumn.pos
        if pos > (prev_pos + 1):
            # need to pad zeros
            hist_dict[0] += (pos - prev_pos - 1)
        prev_pos = pos
        hist_dict[pileupcolumn.n] += 1

def plot_depth_of_coverage(cov_depth_hist):    
    x = sorted(cov_depth_hist.keys())
    y = [cov_depth_hist[i] for i in x]
    plt.plot(x,y)
    plt.savefig('covdepth.png')

if __name__ == '__main__':
    import sys
    logging.basicConfig(level=logging.DEBUG)
    infile = sys.argv[1]
    # open the bam file and read the header
    bamfile = pysam.Samfile(infile, 'rb')
    refs = bamfile.references
    lengths = bamfile.lengths
    # process the file to tabulate coverage depth
    cov_depth_hist = collections.defaultdict(lambda: 0)
    for ref, length in zip(refs, lengths):
        logging.info("Analyzing %s" % (ref))
        depth_of_coverage(bamfile.pileup(ref), cov_depth_hist)
        logging.info("Finished %s" % (ref))

    for cov, n in cov_depth_hist.iteritems():
        print '%d\t%d' % (cov, n)
    plot_depth_of_coverage(cov_depth_hist)
